package RBTree

import (
	"fmt"
	"sync"
)

//定义常量红黑
const (
	RED=true
	BLACK=false
)

var rLock sync.RWMutex

//红黑树节点结构
type rbNode struct{
	Left    *rbNode     //左节点
	Right   *rbNode     //右边节点
	Parent  *rbNode     //父亲节点
	Color   bool        //颜色
	Item                //数据接口
}
//红黑树
type rbTree struct{
	NIL     *rbNode
	Root    *rbNode
	count   uint
	status  bool    //树的状态, true 为正在调整树的平衡
}

//比大小
func  less(x,y Item)bool{
	return x.Less(y)
}

//初始化内存
func NewRBTree()*rbTree {
	return new(rbTree).init()
}
//初始化红黑树
func (rbt *rbTree)init() *rbTree {
	node:=&rbNode{nil,nil,nil,BLACK, nil}
	return &rbTree{node,node,0, false}
}

//获取红黑树元素个数
func (rbt *rbTree)Len()uint{
	return rbt.count
}

//取得红黑树的极大值节点
func  (rbt *rbTree)max(x*rbNode)*rbNode{
	if x==rbt.NIL{
		return rbt.NIL
	}
	for x.Right!=rbt.NIL{
		x=x.Right
	}
	return x
}
//取得红黑树的极小值
func  (rbt *rbTree)min(x*rbNode)*rbNode{
	if x==rbt.NIL{
		return rbt.NIL
	}
	for x.Left!=rbt.NIL{
		x=x.Left
	}
	return x
}


//查找元素, 返回nil 说明没找到
func (rbt *rbTree) Search(item Item) Item {
	if item == nil {
		return rbt.NIL
	}
	tmp := rbt.search(&rbNode{rbt.NIL, rbt.NIL, rbt.NIL, RED, item})
	if tmp != nil {
		return tmp.Item.GetInfo()
	} else {
		return rbt.NIL
	}
}

//搜索红黑树
func (rbt *rbTree) search(x*rbNode)*rbNode{
	for true {
		if rbt.status == false {
			break
		}
	}
	pNode:=rbt.Root         //根节点
	for pNode!=rbt.NIL{
		if less(pNode.Item,x.Item){
			pNode=pNode.Right
		}else if less(x.Item,pNode.Item){
			pNode=pNode.Left
		}else{
			return pNode
		}
	}
	return nil
}
/*
左旋情况
					 80												80
			       /    \										  /	  \
			      40	 120									 60    120
                 /	\			========》以40节点 左旋			/  \
				20	 60										  40    70
                    /  \									 /	\
                   50   70									20   50
 */

func (rbt *rbTree) leftRotate(x*rbNode){
	if  x.Right==rbt.NIL{
		return  //左旋转，逆时针，右孩子不可以为0
	}
	y:=x.Right
	x.Right=y.Left //实现旋转的左旋
	if y.Left!=rbt.NIL{
		y.Left.Parent=x //设定父亲节点
	}
	y.Parent=x.Parent //传递父节点
	if x.Parent==rbt.NIL{
		//根节点
		rbt.Root=y
	}else if x==x.Parent.Left{//x在根节点左边
		x.Parent.Left=y
	}else{//x在根节点右边
		x.Parent.Right=y
	}
	y.Left=x
	x.Parent=y
}

/*
右旋情况
					 80												80
			       /    \										  /	  \
			      60	 120									 40    120
                 /	\			========》以40节点 右旋			/  \
				40	 70										  20    60
               /  \							                 	    / \
              20   50									          50   70
*/

func (rbt *rbTree) rightRotate(x*rbNode){
	if x.Left==nil{
		return //右边旋转，左子树不可以为空
	}
	y:=x.Left
	x.Left=y.Right
	if y.Right!=rbt.NIL{
		y.Right.Parent=x //设置祖先
	}
	y.Parent=x.Parent //y保存x的父亲节点
	if x.Parent==rbt.NIL{
		rbt.Root=y
	}else if x==x.Parent.Left{ //x小于根节点
		x.Parent.Left=y  //父亲节点的孩子是x,改。父亲节点孩子y
	}else{//x大于根节点
		x.Parent.Right=y
	}
	y.Right=x
	x.Parent=y
}

//插入一条数据
func  (rbt *rbTree)Insert(item Item) {
	rbt.status = true
	rLock.Lock()
	if item == nil{
		rLock.Unlock()
		return
	}
	rbt.insert(&rbNode{rbt.NIL,rbt.NIL,rbt.NIL,RED,item})
	rLock.Unlock()
	rbt.status = false
}

//插入
func (rbt *rbTree)insert(z*rbNode)*rbNode{
	//寻找插入位置
	x:=rbt.Root
	y:=rbt.NIL

	for x!=rbt.NIL{
		y=x //备份位置，数据插入x,y之间
		if less(z.Item,x.Item){ //小于
			x=x.Left
		}else if less(x.Item,z.Item){//大于
			x=x.Right
		}else{//相等
			x.Item = z.Item
			return x
		}
	}
	z.Parent=y

	if y==rbt.NIL{
		rbt.Root=z
	}else if less(z.Item,y.Item){
		y.Left=z //小于左边插入
	}else{
		y.Right=z //大于右边插入
	}
	rbt.count++
	rbt.insertFixup(z) //调整平衡
	return z
}

//插入之后，调整平衡
func (rbt*rbTree) insertFixup(z*rbNode){
	for  z.Parent.Color==RED{//一直循环下去，直到根节点
		if z.Parent==z.Parent.Parent.Left{   //父亲节点在爷爷左边
			y:=z.Parent.Parent.Right
			if  y.Color==RED{ //判断大伯节点红色，黑色
				z.Parent.Color=BLACK
				y.Color=BLACK
				z.Parent.Parent.Color=RED
				z=z.Parent.Parent//循环前进
			}else{
				if  z==z.Parent.Right{  //z比父亲小
					z=z.Parent
					rbt.leftRotate(z)//左旋
				}else{//z比父亲大
					z.Parent.Color=BLACK
					z.Parent.Parent.Color=RED
					rbt.rightRotate(z.Parent.Parent)
				}
			}
		}else{// //父亲节点在爷爷右边
			y:=z.Parent.Parent.Left//叔叔节点
			if  y.Color==RED{ //判断大伯节点红色，黑色
				z.Parent.Color=BLACK
				y.Color=BLACK
				z.Parent.Parent.Color=RED
				z=z.Parent.Parent //循环前进
			}else{
				if z==z.Parent.Left{
					z=z.Parent
					rbt.rightRotate(z)
				}else{
					z.Parent.Color=BLACK
					z.Parent.Parent.Color=RED
					rbt.leftRotate(z.Parent.Parent)
				}
			}
		}
	}
	rbt.Root.Color=BLACK
}


//获取树的高度
func (rbt *rbTree) GetDepth() int {
	return rbt.getDepth(rbt.Root)
}

func (rbt *rbTree) getDepth(root *rbNode ) int {
	if root == rbt.NIL {
		return 0
	}
	if root.Left == rbt.NIL || root.Right == rbt.NIL {
		return 1
	}
	left := rbt.getDepth(root.Left)
	right := rbt.getDepth(root.Right)
	if left > right {
		return left + 1
	} else {
		return right + 1
	}
}

//近似查找
func (rbt *rbTree) SearchLe(item Item) Item {
	for true {
		if rbt.status == false {
			break
		}
	}
	if item == rbt.NIL {
		return rbt.NIL
	}
	return rbt.searchLe(&rbNode{rbt.NIL, rbt.NIL, rbt.NIL, BLACK, item}).Item.GetInfo()
}

func (rbt *rbTree)searchLe(x*rbNode)*rbNode{
	p:=rbt.Root //根节点
	n:=p //备份根节点
	for n!=rbt.NIL{
		if less(n.Item,x.Item){
			p=n
			n=n.Right //大于
		}else if less(x.Item,n.Item){
			p=n
			n=n.Left//小于
		}else{
			return n
		}
	}
	if less(p.Item,x.Item){
		return p
	}
	p=rbt.deSuccessor(p)
	return p
}
func (rbt *rbTree)successor(x*rbNode)*rbNode{
	if  x==rbt.NIL{
		return rbt.NIL
	}
	if x.Right!=rbt.NIL{
		return rbt.min(x.Right) //取得右边最小
	}
	y:=x.Parent
	for y!=rbt.NIL && x==y.Right{
		x=y
		y=y.Parent
	}
	return y
}
func (rbt *rbTree)deSuccessor(x*rbNode)*rbNode{
	if  x==rbt.NIL{
		return rbt.NIL
	}
	if x.Left!=rbt.NIL{
		return rbt.max(x.Left) //取得左边最大
	}
	y:=x.Parent
	for y!=rbt.NIL && x==y.Left{
		x=y
		y=y.Parent
	}
	return y
}

//前序遍历
func (rbt *rbTree) PrePrintAll() {
	rbt.prePrintAll(rbt.Root)
}

func (rbt *rbTree) prePrintAll(n *rbNode) {
	if n == rbt.NIL {
		return
	}
	fmt.Println(n.Item.GetInfo())
	rbt.prePrintAll(n.Left)
	rbt.prePrintAll(n.Right)
}

//中序遍历
func (rbt *rbTree) InPrintAll() {
	rbt.inPrintAll(rbt.Root)
}

func (rbt *rbTree) inPrintAll(n *rbNode) {
	if n == rbt.NIL {
		return
	}
	rbt.inPrintAll(n.Left)
	fmt.Println(n.Item.GetInfo())
	rbt.inPrintAll(n.Right)
}

//后序遍历
func (rbt *rbTree) PostPrintAll() {
	rbt.postPrintAll(rbt.Root)
}

func (rbt *rbTree) postPrintAll(n *rbNode) {
	if n == rbt.NIL {
		return
	}
	rbt.postPrintAll(n.Left)
	rbt.postPrintAll(n.Right)
	fmt.Println(n.Item.GetInfo())
}

//删除数据
func  (rbt *rbTree)Delete(item Item) Item {
	for true {
		if rbt.status == false {
			break
		}
	}
	if item == nil {
		return nil
	}
	return rbt.delete(&rbNode{rbt.NIL,rbt.NIL,rbt.NIL,RED,item}).Item
}

func (rbt *rbTree)delete(key*rbNode)*rbNode{
	z:=rbt.search(key)//寻找要删除的节点
	if z==rbt.NIL{
		return rbt.NIL//无需删除
	}
	//新建节点下，x,y备份,夹逼
	var x *rbNode
	var y *rbNode

	//节点
	ret:=&rbNode{rbt.NIL,rbt.NIL,rbt.NIL,z.Color,z.Item}

	if z.Left==rbt.NIL ||z.Right==rbt.NIL{
		y=z //单节点，y,z重合
	}else{
		y=rbt.successor(z)//找到最接近的右边最小
	}

	if y.Left!=rbt.NIL{
		x=y.Left
	}else{
		x=y.Right
	}
	x.Parent=y.Parent

	if y.Parent==rbt.NIL{
		rbt.Root=x
	}else if y==y.Parent.Left{
		y.Parent.Left=x
	}else{
		y.Parent.Right=x
	}
	if y!=z{
		z.Item=y.Item
	}
	if y.Color==BLACK{
		rbt.deleteFixup(x)      //调整平衡
	}
	rbt.count--
	return ret
}

//删除时的红黑修复需要考虑四种情况(下面的“X”指取代了被删节点位置的新节点)
// (1) X的兄弟节点是红色的：
//        |                              |
//       1●                             3●
//       / \             --------\      / \
// X-> 2●   ○3 <-brother --------/    1○   ●5
//         / \                        / \
//       4●   ●5                X-> 2●   ●4
//
// (2) X的兄弟节点是黑色的，而且兄弟节点的两个孩子都是黑色的：
//        |                              |
//       1○                         X-> 1○
//       / \             --------\      / \
// X-> 2●   ●3 <-brother --------/    2●   ○3
//         / \                            / \
//       4●   ●5                        4●   ●5
//
// (3) X的兄弟节点是黑色的，兄弟的左孩子是红色的，右孩子是黑色的：
//        |                              |
//       1○                             1○
//       / \             --------\      / \
// X-> 2●   ●3 <-brother --------/ X->2●   ●4
//         / \                              \
//       4○   ●5                             ○3
//                                            \
//                                             ●5
//
// (4) X的兄弟节点是黑色的，兄弟的右孩子是红色的：
//        |                              |
//       1○                             3○   X->root and loop while end
//       / \             --------\      / \
// X-> 2●   ●3 <-brother --------/    1●   ●5
//         / \                        / \
//       4○   ○5                    2●   ○4
//
//以上是兄弟节点在X右边时的情况，在X左边是取相反即可!
//删除节点, 保持平衡
func (rbt *rbTree)deleteFixup(x*rbNode){
	for x!=rbt.Root && x.Color==BLACK{
		if x==x.Parent.Left{  //x在左边
			w:=x.Parent.Right//哥哥节点
			if  w.Color==RED{ //左边旋转
				w.Color=BLACK
				x.Parent.Color=RED
				rbt.leftRotate(x.Parent)
				w=x.Parent.Right //循环步骤
			}
			if w.Left.Color==BLACK  && w.Right.Color==BLACK{
				w.Color=RED
				x=x.Parent //循环条件
			}else{
				if w.Right.Color==BLACK{
					w.Left.Color=BLACK
					w.Color=RED
					rbt.rightRotate(w)//右旋转
					w=x.Parent.Right //循环条件
				}
				w.Color=x.Parent.Color
				x.Parent.Color=BLACK
				w.Right.Color=BLACK
				rbt.leftRotate(x.Parent)
				x=rbt.Root
			}
		}else{ //x在右边
			w:=x.Parent.Left//左边节点
			if  w.Color==RED{ //左旋
				w.Color=BLACK
				x.Parent.Color=RED
				rbt.rightRotate(x.Parent)
				w=x.Parent.Right //循环步骤
			}
			if w.Left.Color==BLACK  && w.Right.Color==BLACK{

				w.Color=RED
				x=x.Parent //循环条件
			}else{
				if w.Right.Color==BLACK{
					w.Left.Color=BLACK
					w.Color=RED
					rbt.leftRotate(w)//右旋转
					w=x.Parent.Left//循环条件
				}
				w.Color=x.Parent.Color
				x.Parent.Color=BLACK
				w.Right.Color=BLACK
				rbt.rightRotate(x.Parent)
				x=rbt.Root
			}
		}
	}
	x.Color=BLACK//循环到最后根节点，黑色
}